package com.zh.mediaservice.AiBigModel;

import cn.hutool.core.util.StrUtil;
import com.zh.mediacommon.constant.BigModelSocketConstant;
import com.zh.mediacommon.ecxeption.AIModelException;
import com.zh.mediapojo.dto.MsgDTO;


import com.zh.mediaservice.component.XfXhStreamClient;
import com.zh.mediaservice.config.XfXhConfig;
import com.zh.mediaservice.listener.XfXhWebSocketListener;
import lombok.extern.slf4j.Slf4j;
import okhttp3.WebSocket;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Collections;

@Component
@Slf4j
public class SparkUltra {

    @Resource
    private XfXhStreamClient xfXhStreamClient;

    @Resource
    private XfXhConfig xfXhConfig;
    public  String getResponse(String question, Integer uid) {
        // 如果是无效字符串，则不对大模型进行请求
        if (StrUtil.isBlank(question)) {
           throw new AIModelException(BigModelSocketConstant.INVALID_QUESTION);
        }
        // 获取连接令牌
        if (!xfXhStreamClient.operateToken(XfXhStreamClient.GET_TOKEN_STATUS)) {
            throw new AIModelException(BigModelSocketConstant.OVERLOAD);
        }

        // 创建消息对象
        MsgDTO msgDTO = MsgDTO.createUserMsg(question);
        // 创建监听器
        XfXhWebSocketListener listener = new XfXhWebSocketListener();
        // 发送问题给大模型，生成 websocket 连接
        WebSocket webSocket = xfXhStreamClient.sendMsg(uid.toString(), Collections.singletonList(msgDTO), listener);
        if (webSocket == null) {
            // 归还令牌
            xfXhStreamClient.operateToken(XfXhStreamClient.BACK_TOKEN_STATUS);
            throw new AIModelException(BigModelSocketConstant.ERROR);
        }
        try {
            int count = 0;
            // 为了避免死循环，设置循环次数来定义超时时长
            int maxCount = xfXhConfig.getMaxResponseTime();
            while (count <= maxCount) {
                Thread.sleep(200);
                if (listener.isWsCloseFlag()) {
                    break;
                }
                count++;
            }
            if (count > maxCount) {
                throw new AIModelException(BigModelSocketConstant.INTERNET_ERROR);
            }
            // 响应大模型的答案

            return listener.getAnswer().toString();
        } catch (InterruptedException e) {
            log.error("错误：" + e.getMessage());
            throw new AIModelException(BigModelSocketConstant.ERROR);
        } finally {
            // 关闭 websocket 连接
            webSocket.close(1000, "");
            // 归还令牌
            xfXhStreamClient.operateToken(XfXhStreamClient.BACK_TOKEN_STATUS);
        }
    }

    }
